from PIL import Image
import numpy as np
import os
import math

def _check_blocks(img_array, color_threshold, color_ratio_threshold, saturation_threshold=0.15, hist_diff_threshold=0.2):
    """
    辅助函数：对单个图像块进行颜色检查（多维度方法）。
    
    :param saturation_threshold: 平均饱和度阈值，高于此值判定为彩色。
    :param hist_diff_threshold: 颜色通道直方图差异阈值，高于此值判定为彩色。
    """
    height, width, _ = img_array.shape
    total_block_pixels = height * width
    if total_block_pixels == 0:
        return True

    # --- 方法1: 灰度像素比例检查 (基础方法) ---
    r, g, b = img_array[:,:,0], img_array[:,:,1], img_array[:,:,2]
    diff_rg = np.abs(r.astype(int) - g.astype(int))
    diff_rb = np.abs(r.astype(int) - b.astype(int))
    diff_gb = np.abs(g.astype(int) - b.astype(int))
    gray_mask = (diff_rg < color_threshold) & (diff_rb < color_threshold) & (diff_gb < color_threshold)
    gray_ratio = np.sum(gray_mask) / total_block_pixels

    if gray_ratio < color_ratio_threshold:
        print(f"块判定为彩色 (方法1: 灰度像素占比过低: {gray_ratio:.2f})")
        return False

    # --- 方法2: 平均饱和度检查 (高效且准确) ---
    # 将 NumPy 数组转回 PIL Image 以便转换到 HSV
    block_img = Image.fromarray(img_array)
    hsv_img = block_img.convert('HSV')
    hsv_array = np.array(hsv_img)
    # S (Saturation) 通道在索引1，其值范围是 0-255
    avg_saturation = np.mean(hsv_array[:, :, 1]) / 255.0  # 归一化到 0-1

    if avg_saturation > saturation_threshold:
        print(f"块判定为彩色 (方法2: 平均饱和度过高: {avg_saturation:.2f})")
        return False

    # --- 方法3: 颜色直方图差异检查 ---
    # 计算R, G, B三个通道的直方图
    hist_r, _ = np.histogram(r, bins=256, range=(0, 256))
    hist_g, _ = np.histogram(g, bins=256, range=(0, 256))
    hist_b, _ = np.histogram(b, bins=256, range=(0, 256))

    # 归一化直方图
    hist_r = hist_r / total_block_pixels
    hist_g = hist_g / total_block_pixels
    hist_b = hist_b / total_block_pixels

    # 计算直方图之间的差异
    diff_rg = np.sum(np.abs(hist_r - hist_g))
    diff_rb = np.sum(np.abs(hist_r - hist_b))
    
    # 如果两个差异都很大，说明颜色通道不平衡
    if diff_rg > hist_diff_threshold and diff_rb > hist_diff_threshold:
        print(f"块判定为彩色 (方法3: 颜色直方图差异过大: R-G={diff_rg:.2f}, R-B={diff_rb:.2f})")
        return False
        
    return True


def is_black_white_comic(image_path, color_threshold=50, color_ratio_threshold=0.5, min_block_pixels=50000, black_threshold=10,saturation_threshold=0.15, hist_diff_threshold=0.2, white_threshold=245, std_dev_threshold=20):
    """
    判断图片是否为黑白漫画（横向和纵向分别检查）
    """
    try:
        img = Image.open(image_path).convert('RGB')
        img_array = np.array(img)
        
        # 排除纯黑/纯白图片
        mean_brightness = np.mean(img_array)
        std_dev = np.std(img_array)
        if (mean_brightness < black_threshold or mean_brightness > white_threshold) and std_dev < std_dev_threshold:
            print(f"图片被判断为纯黑/纯白或颜色单一，跳过。")
            return False

        height, width, _ = img_array.shape
        total_pixels = height * width
        print(f"图片尺寸: {width}x{height}, 总像素: {total_pixels}")
        # 如果图片总像素都小于最小块要求，则不分块，直接检查整张图
        if total_pixels < min_block_pixels:
            print("图片总像素小于最小分块要求，将整图作为单块检查...")
            return _check_blocks(img_array, color_threshold, color_ratio_threshold,saturation_threshold, hist_diff_threshold)

        # --- 1. 横向条带检查 ---
        # 每个横向条带的宽度是固定的 (width)，要满足最小像素要求，其高度至少应为 min_block_pixels / width
        min_h_strip = math.ceil(min_block_pixels / width)
        num_h_blocks = math.floor(height / min_h_strip) if min_h_strip > 0 else 1
        num_h_blocks = max(1, num_h_blocks) # 至少分为1块

        block_h = height // num_h_blocks
        if block_h > 0:
            print(f"进行横向条带检查 (共 {num_h_blocks} 块)...")
            for i in range(num_h_blocks):
                block = img_array[i*block_h:(i+1)*block_h, :]
                if not _check_blocks(block, color_threshold, color_ratio_threshold,saturation_threshold, hist_diff_threshold):
                    return False
        
        # --- 2. 纵向条带检查 ---
        # 每个纵向条带的高度是固定的 (height)，要满足最小像素要求，其宽度至少应为 min_block_pixels / height
        min_w_strip = math.ceil(min_block_pixels / height)
        num_v_blocks = math.floor(width / min_w_strip) if min_w_strip > 0 else 1
        num_v_blocks = max(1, num_v_blocks) # 至少分为1块

        block_w = width // num_v_blocks
        if block_w > 0:
            print(f"进行纵向条带检查 (共 {num_v_blocks} 块)...")
            for i in range(num_v_blocks):
                block = img_array[:, i*block_w:(i+1)*block_w]
                if not _check_blocks(block, color_threshold, color_ratio_threshold):
                    return False

        return True
    except Exception as e:
        print(f"处理图片时出错: {e}")
        return False

def is_color_comic(image_path, color_threshold=50, color_ratio_threshold=0.5, min_block_pixels=50000,
                   black_threshold=10, white_threshold=245, std_dev_threshold=20,
                   saturation_threshold=0.15, hist_diff_threshold=0.2, color_block_ratio=1/3):
    """
    判断图片是否为彩色漫画（分块检测，超过1/3块为彩色则判定为彩色）
    """
    try:
        img = Image.open(image_path).convert('RGB')
        img_array = np.array(img)

        # 排除纯黑/纯白图片
        mean_brightness = np.mean(img_array)
        std_dev = np.std(img_array)
        if (mean_brightness < black_threshold or mean_brightness > white_threshold) and std_dev < std_dev_threshold:
            return False

        height, width, _ = img_array.shape
        total_pixels = height * width
        if total_pixels < min_block_pixels:
            blocks = [img_array]
        else:
            # 横向分块
            min_h_strip = math.ceil(min_block_pixels / width)
            num_h_blocks = max(1, math.floor(height / min_h_strip))
            block_h = height // num_h_blocks
            blocks = [img_array[i*block_h:(i+1)*block_h, :] for i in range(num_h_blocks)]
            # 纵向分块
            min_w_strip = math.ceil(min_block_pixels / height)
            num_v_blocks = max(1, math.floor(width / min_w_strip))
            block_w = width // num_v_blocks
            blocks += [img_array[:, i*block_w:(i+1)*block_w] for i in range(num_v_blocks)]

        color_block_count = 0
        valid_block_count = 0
        for block in blocks:
            if block.size == 0:
                continue
            valid_block_count += 1
            # 判定该块是否为彩色
            if not _check_blocks(block, color_threshold, color_ratio_threshold,
                                 saturation_threshold, hist_diff_threshold):
                color_block_count += 1

        if valid_block_count == 0:
            return False

        # 超过1/3块为彩色则判定为彩色漫画
        if color_block_count / valid_block_count >= color_block_ratio:
            return True
        return False
    except Exception as e:
        print(f"处理图片时出错: {e}")
        return False

# 使用示例
if __name__ == "__main__":
    import time
    rootdir = "E:\\software\\jmcomic_v1.2.6_windows_x64\\jmcomic\\cache\\book\\1052306\\1"
    not_color=[]
    color=[]
    for imgname in os.listdir(rootdir):
        print("----------")
        print(imgname)
        start_time = time.time()
        result = is_black_white_comic(os.path.join(rootdir,imgname),color_ratio_threshold=0.5)
        end_time = time.time()
        print(f"该图片是黑白漫画: {result}")
        print(f"处理时间: {end_time - start_time:.2f}秒")
        if not result:
            color.append(imgname)
        else:
            not_color.append(imgname)
    print("非黑白漫画图片列表：", color)
    print("黑白漫画图片列表：", not_color)


